package com.means.polymerize.core.terminal.common.util;

import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.ssh.ChannelType;
import cn.hutool.extra.ssh.JschUtil;
import com.jcraft.jsch.Channel;
import com.jcraft.jsch.Session;
import com.means.polymerize.common.constants.Global;
import com.means.polymerize.core.terminal.common.lang.ChannelStorageContainer;
import com.means.polymerize.core.terminal.common.lang.SessionStorageContainer;
import com.means.polymerize.core.terminal.modal.entity.TerminalConnectInfoReceive;
import com.means.polymerize.core.terminal.service.TerminalService;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.atomic.LongAdder;

/**
 * ssh工具
 *
 * @author xuhaifeng
 * @date 2021.05.09
 */
@Slf4j
@Component
public class SshUtil {
    private static TerminalService service;

    /**
     * Connect ssh.
     *
     * @param token            the token
     * @param sessionContainer the session container
     * @param receive          the receive
     */
    @SneakyThrows
    public static void connectSsh(String token, TerminalConnectInfoReceive receive) {
        /* 激活会话，5秒超时 */
        final AtomicReference<Session> session = new AtomicReference<>();
        try {
            session.set(JschUtil.openSession(receive.getHost(), receive.getPort(), receive.getUser(), receive.getSecret(), 5000));
        } catch (Exception e) {
            service.closeSessionConnection(token);
            return;
        }
        /* 开启shell通道，5秒超时 */
        final Channel channel = JschUtil.openChannel(session.get(), ChannelType.SHELL, 5000);
        /* 存储通道 */
        ChannelStorageContainer.setChannel(token, channel);
        /* 发送消息 */
        sendSshCommand(channel, Global.NOT);
        /* 读取终端返回的信息流 */
        try (InputStream inputStream = channel.getInputStream()) {
            /* 配置每次块大小 */
            final byte[] buffer = new byte[1024];
            /* 当前响应内容总数 */
            final LongAdder read = new LongAdder();
            /* 如果没有数据来，线程会一直阻塞在这个地方等待数据 */
            while (read.sum() >= Global.NUM_I_0) {
                read.add(inputStream.read(buffer, Global.NUM_I_0, buffer.length) - read.sum());
                if (read.sum() > Global.NUM_I_0) {
                    log.info("当前响应内容总量：%s".formatted(read));
                    sendMessage(token, Arrays.copyOfRange(buffer, Global.NUM_I_0, read.intValue()));
                }
            }
        } finally {
            session.get().disconnect();
            channel.disconnect();
        }
    }

    /**
     * 发送ssh指令
     *
     * @param channel channel
     * @param command command
     */
    @SneakyThrows
    public static void sendSshCommand(Channel channel, String command) {
        if (Objects.nonNull(channel)) {
            final OutputStream outputStream = channel.getOutputStream();
            outputStream.write(command.getBytes());
            outputStream.flush();
        }
    }

    /**
     * 发送消息
     *
     * @param buffer buffer
     */
    @SneakyThrows
    private static void sendMessage(String token, byte[] buffer) {
        SessionStorageContainer.getSessions(token).sendText(StrUtil.utf8Str(buffer));
    }

    @Autowired
    public void setSshUtil(TerminalService service) {
        SshUtil.service = service;
    }
}
